import numpy as np
import scipy.linalg
import os
import sys
from math import pi
from matplotlib.pylab import *

if __name__=='__main__':

    t = np.loadtxt('time.dat');
    tau = np.loadtxt('tau.dat');
    rho0 = np.loadtxt("rho.dat")
    H  = np.loadtxt("hamiltonian_real.dat")
    H += np.loadtxt("hamiltonian_imag.dat")
    
    a =  1.0
    c = -1.0
    b = np.sqrt(8.0)

    E,U=scipy.linalg.eig(H)

    print H
    print E
    Uh = U.T.conj()

    rho = rho0.copy()
    f = np.zeros((t.shape[0],tau.shape[0]),dtype=np.complex)
    f[0,0] = rho[0,0] 

    rhoi_c = np.zeros((rho.shape[0],rho.shape[1],tau.shape[0]),dtype=np.complex)
    rhoi_n = rhoi_c.copy()
    rhoi_c[:,:,0] = rho
    I = np.eye(rho.shape[0])
    for j in range(1,t.shape[0]):
        dt = t[j]-t[j-1]
        u = scipy.linalg.matfuncs.expm2(-1j*H*dt)        
        rho = np.dot(u,np.dot(rho,u.T.conj()))

        rhoi_n[:,:,0] = rho

        for i in range(1,tau.shape[0]):
            dt = tau[i]-tau[i-1]
            #u = scipy.linalg.matfuncs.expm2(-1j*H*dt)
            # The following implementation matches that used in propagator.cpp
            # where only the diagonal terms are propagated using exponentiated matrix
            g  = np.dot((I - 1j*H*dt/2),rhoi_c[:,:,i-1])
            g  = scipy.linalg.solve(I + 1j*H*dt/2,g)
            rhoi_n[:,:,i] = g

        
        f[j,:] = rhoi_n[0,0,:]


        rhoi_c = rhoi_n.copy()


    # Solution by the qk code
    fr0=np.loadtxt('sol_real.dat');
    fi0=np.loadtxt('sol_imag.dat');

    fr1=np.real(f);
    fi1=np.imag(f);

    print "|fqk - fpy|=",np.sqrt(np.sum(np.abs(fr0+1j*fi0-f)**2))

    if (1):

        aps = {'params': { 'axes.labelsize':  10,
                           'text.fontsize':   10,
                           'legend.fontsize': 10,
                           'xtick.labelsize': 8,
                           'ytick.labelsize': 8,
                           'text.usetex': True,
                           'figure.figsize': (4,7)},
                           'axes': [0.125,0.2,0.95-0.125,0.95-0.2]}
    
        rcParams.update(aps['params'])    
        rcParams['text.dvipnghack']=True
        rcParams['font.family']='serif'
        rcParams['font.serif']='Computer Modern Roman'

        clf();
        subplot(2,1,1);
        plot(t,fr0[:,[0,-1]])
        plot(t,fr1[:,[0,-1]],'-.');
        legend(["qk 00,tau=0","qk 00,tau=end","py 00,tau=0","py 00,tau=end"])
        ylabel("real")
        subplot(2,1,2);
        plot(t,fi0[:,[0,-1]]);
        plot(t,fi1[:,[0,-1]],'-.');
        xlabel("time")
        ylabel("imaginary")
        legend(["qk 00,tau=0","qk 00,tau=end","py 00,tau=0","py 00,tau=end"])
    
        savefig("test2x2tau.pdf")

